{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import json\n",
    "from PIL import Image, ImageDraw\n",
    "import os\n",
    "import cv2\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import shutil\n",
    "import random"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The purpose of this script is to explore images/annotations of the WIDER dataset.  \n",
    "Also it converts annotations into json format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# first run this script for this images:\n",
    "IMAGES_DIR = '/home/gpu2/hdd/dan/WIDER/WIDER_train/images/'\n",
    "BOXES_PATH = '/home/gpu2/hdd/dan/WIDER/wider_face_split/wider_face_train_bbx_gt.txt'\n",
    "RESULT_DIR = '/home/gpu2/hdd/dan/WIDER/train/'\n",
    "\n",
    "# then run for this images:\n",
    "# IMAGES_DIR = '/home/gpu2/hdd/dan/WIDER/WIDER_val/images/'\n",
    "# BOXES_PATH = '/home/gpu2/hdd/dan/WIDER/wider_face_split/wider_face_val_bbx_gt.txt'\n",
    "# RESULT_DIR = '/home/gpu2/hdd/dan/WIDER/train_part2/'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Read data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# collect paths to all images\n",
    "\n",
    "all_paths = []\n",
    "for path, subdirs, files in tqdm(os.walk(IMAGES_DIR)):\n",
    "    for name in files:\n",
    "        all_paths.append(os.path.join(path, name))\n",
    "        \n",
    "metadata = pd.DataFrame(all_paths, columns=['full_path'])\n",
    "\n",
    "# strip root folder\n",
    "metadata['path'] = metadata.full_path.apply(lambda x: os.path.relpath(x, IMAGES_DIR))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# see all unique endings\n",
    "metadata.path.apply(lambda x: x.split('.')[-1]).unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# number of images\n",
    "len(metadata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# read annotations\n",
    "with open(BOXES_PATH, 'r') as f:\n",
    "    content = f.readlines()\n",
    "    content = [s.strip() for s in content]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# split annotations by image\n",
    "boxes = {}\n",
    "num_lines = len(content)\n",
    "i = 0\n",
    "name = None\n",
    "\n",
    "while i < num_lines:\n",
    "    s = content[i]\n",
    "    if s.endswith('.jpg'):\n",
    "        if name is not None:\n",
    "            assert len(boxes[name]) == num_boxes\n",
    "        name = s\n",
    "        boxes[name] = []\n",
    "        i += 1\n",
    "        num_boxes = int(content[i])\n",
    "        i += 1\n",
    "    else:\n",
    "        xmin, ymin, w, h = s.split(' ')[:4]\n",
    "        xmin, ymin, w, h = int(xmin), int(ymin), int(w), int(h)\n",
    "        if h <= 0 or w <= 0:\n",
    "            print(name)  \n",
    "            # some boxes are weird!\n",
    "            # so i don't use them\n",
    "            num_boxes -= 1\n",
    "        else:\n",
    "            boxes[name].append((xmin, ymin, w, h))\n",
    "        i += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check that all images have bounding boxes\n",
    "assert metadata.path.apply(lambda x: x in boxes).all()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Show some bounding boxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_boxes_on_image(path, boxes):\n",
    "\n",
    "    image = Image.open(path)\n",
    "    draw = ImageDraw.Draw(image, 'RGBA')\n",
    "    width, height = image.size\n",
    "\n",
    "    for b in boxes:\n",
    "        xmin, ymin, w, h = b\n",
    "        xmax, ymax = xmin + w, ymin + h\n",
    "\n",
    "        fill = (255, 255, 255, 45)\n",
    "        outline = 'red'\n",
    "        draw.rectangle(\n",
    "            [(xmin, ymin), (xmax, ymax)],\n",
    "            fill=fill, outline=outline\n",
    "        )\n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "i = random.randint(0, len(metadata) - 1)  # choose a random image\n",
    "some_boxes = boxes[metadata.path[i]]\n",
    "draw_boxes_on_image(metadata.full_path[i], some_boxes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Convert"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_annotation(path, width, height):\n",
    "    name = path.split('/')[-1]\n",
    "    annotation = {\n",
    "      \"filename\": name,\n",
    "      \"size\": {\"depth\": 3, \"width\": width, \"height\": height}\n",
    "    }\n",
    "    objects = []\n",
    "    for b in boxes[path]:\n",
    "        xmin, ymin, w, h = b\n",
    "        xmax, ymax = xmin + w, ymin + h\n",
    "        objects.append({\n",
    "            \"bndbox\": {\"ymin\": ymin, \"ymax\": ymax, \"xmax\": xmax, \"xmin\": xmin}, \n",
    "            \"name\": \"face\"\n",
    "        })\n",
    "    annotation[\"object\"] = objects\n",
    "    return annotation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a folder for the converted dataset\n",
    "shutil.rmtree(RESULT_DIR, ignore_errors=True)\n",
    "os.mkdir(RESULT_DIR)\n",
    "os.mkdir(os.path.join(RESULT_DIR, 'images'))\n",
    "os.mkdir(os.path.join(RESULT_DIR, 'annotations'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for T in tqdm(metadata.itertuples()):\n",
    "    \n",
    "    # get width and height of an image\n",
    "    image = cv2.imread(T.full_path)\n",
    "    h, w, c = image.shape\n",
    "    assert c == 3\n",
    "    \n",
    "    # name of the image\n",
    "    name = T.path.split('/')[-1]\n",
    "    assert name.endswith('.jpg')\n",
    "\n",
    "    # copy the image\n",
    "    shutil.copy(T.full_path, os.path.join(RESULT_DIR, 'images', name))\n",
    "    \n",
    "    # save annotation for it\n",
    "    d = get_annotation(T.path, w, h)\n",
    "    json_name = name[:-4] + '.json'\n",
    "    json.dump(d, open(os.path.join(RESULT_DIR, 'annotations', json_name), 'w')) "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
